import torch
from pytorch3d.loss import chamfer_distance

if __name__ == '__main__':

    torch.cuda.set_device(0)

    pd_x = torch.randn([32,6890,3]).cuda()
    pd_y = torch.randn_like(pd_x)

    torch3d_chamferDist, _ = chamfer_distance(pd_x,pd_y)
    torch3d_chamferDist_psum, _ = chamfer_distance(pd_x, pd_y, point_reduction='sum')
    torch3d_chamferDist_pmean, _ = chamfer_distance(pd_x, pd_y, point_reduction='mean')
    torch3d_chamferDist_bsum, _ = chamfer_distance(pd_x, pd_y, batch_reduction='sum')
    torch3d_chamferDist_bmean, _ = chamfer_distance(pd_x, pd_y, batch_reduction='mean')
    torch3d_chamferDist_bpsum, _ = chamfer_distance(pd_x, pd_y,
                                                     batch_reduction='sum', point_reduction='sum')


    print("Done")
